// Copyright 2014 Timothy Edward Baldwin

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#define __GNU_SOURCE
#include "navalis.h"

#include <unistd.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <netdb.h>
#include <popt.h>

#include <atomic>
#include <fstream>
#include <iostream>
#include <thread>
using std::cout;
using std::endl;

#include <nettle/md5.h>

const int token_size = 8;

int sock;
int remote_sock;
int tun;
int rand_fd;
uint32_t anycast_address;
std::atomic_bool current_key;

struct sockaddr_in relay_saddr;
struct ipv6_addr prefix_6a44;

struct key {
  u_int8_t key[32];
} keys[2];

void generate_token(uint8_t *token, const ipv6_addr &client, const ipv6_addr &native, const struct key &key) {
  struct md5_ctx ctx;
  md5_init(&ctx);
  md5_update(&ctx, sizeof(key.key), key.key);
  md5_update(&ctx, 16, reinterpret_cast<const uint8_t *>(&native));
  md5_update(&ctx, 12, reinterpret_cast<const uint8_t *>(&client));
  md5_digest(&ctx, token_size, token);
}

inline bool check_token(const uint8_t *token1, const ipv6_addr &client, const ipv6_addr &native) {
  uint8_t token2[token_size];
  generate_token(token2, client, native, keys[0]);
  if (memcmp(token1, token2, token_size) == 0) return true;
  generate_token(token2, client, native, keys[1]);
  return memcmp(token1, token2, token_size) == 0;
}

void recv_ipv6() {
  while (true) {
    udp_packet packet;
    ssize_t size = read(tun, &packet, sizeof(packet));
    if (packet.v6.version == 6
        && size == packet.v6.size()
        && packet.v6.destination.mapped_port != PORT_6a44
        && packet.v6.destination.mapped_port != htons(3653) // TSP
        && (packet.v6.source.prefix.x != htonl(0x20010000)
            || packet.v6.source.private_ipv4 != ~anycast_address)) {

      // Add our IPv4 address and token, so that
      // client may send traffic via this relay
      if (remote_sock && size <= (1280 - 4 - token_size)) {
        memcpy(u8_cast(&packet) + size, &relay_saddr.sin_addr, 4);
        generate_token(u8_cast(&packet) + size + 4,
                       packet.v6.destination,
                       packet.v6.source,
                       keys[current_key]);
        size += 4 + token_size;
      }

      // Send the packet to client
      struct sockaddr_in client;
      client.sin_family = AF_INET;
      client.sin_addr.s_addr = packet.v6.destination.public_ipv4;
      client.sin_port = packet.v6.destination.mapped_port;
      sendto(sock, &packet, size, 0, sa_cast(&client), sizeof(client));
    }
  }
}


void recv_ipv4() {
  while (true) {
    udp_packet packet;
    struct sockaddr_in client;
    socklen_t len = sizeof(client);
    ssize_t size = recvfrom(sock, &packet, sizeof(packet), 0, sa_cast(&client), &len);
    if (size < 40) {
      if (size >= 20 && !packet.bubble.address.prefix) goto send_bubble;
    } else if (packet.v6.version == 6 && size >= packet.v6.size()) {
      if (packet.v6.source.mapped_port == client.sin_port
          && packet.v6.source.public_ipv4 == client.sin_addr.s_addr
          && packet.v6.source.prefix == prefix_6a44.prefix) {
        write(tun, &packet.v6, packet.v6.size());
      } else {
        packet.bubble.nounce = 0;
send_bubble:
        memcpy(packet.bubble.prefix, &prefix_6a44, 6);
        packet.bubble.address.public_ipv4 = client.sin_addr.s_addr;
        packet.bubble.address.mapped_port = client.sin_port;
        sendto(sock, &packet.bubble, sizeof(packet.bubble), 0, sa_cast(&client), sizeof(client));
      }
    }
  }
}

void recv_ipv4_direct() {
  while (true) {
    udp_packet packet;
    struct sockaddr_in client;
    socklen_t len = sizeof(client);
    ssize_t size = recvfrom(remote_sock, &packet, sizeof(packet), 0, sa_cast(&client), &len);
    if (packet.v6.version == 6
        && size >= packet.v6.size() + token_size
        && packet.v6.source.public_ipv4 == client.sin_addr.s_addr
        && packet.v6.source.prefix == prefix_6a44.prefix
        && check_token(u8_cast(&packet) + packet.v6.size(),
                       packet.v6.source, packet.v6.destination)) {
      write(tun, &packet.v6, packet.v6.size());
    }
  }
}

namespace args {

const char *pidfile;
const char *ipv4;
const char *remote_ipv4;
const char *prefix;
const char *interface = "6a44_relay";
const char *interface_up;
int no_fork;

struct poptOption options[] = {
{"pidfile", 0, POPT_ARG_STRING, &pidfile, 0, "PID File", nullptr},
{"ipv4", '4', POPT_ARG_STRING, &ipv4, 0, "IPv4 address (required)", nullptr},
{"remote-ipv4", 'r', POPT_ARG_STRING, &remote_ipv4, 0, "IPv4 address for remote clients", nullptr},
{"prefix", '6', POPT_ARG_STRING, &prefix, 0, "IPv6 /48 prefix (required)", nullptr},
{"interface-name", 'i', POPT_ARG_STRING, &interface, 0, "Interface name", nullptr},
{"interface-up", 'u', POPT_ARG_STRING, &interface_up, 0, "Interface up command", nullptr},
{"no-fork", 'f', POPT_ARG_VAL, &no_fork, 1, "Don't fork into background", nullptr},
POPT_AUTOHELP
POPT_TABLEEND
};
}

int main(int argc, const char **argv) {

  // Read command line arguments
  poptContext opts = poptGetContext(nullptr, argc, argv, args::options, 0);
  poptGetNextOpt(opts);
  poptFreeContext(opts);

  inet_pton(AF_INET6, args::prefix, &prefix_6a44);

  // Open random number device file
  rand_fd = open("/dev/urandom", O_RDONLY | O_CLOEXEC);

  // Initialise keys
  set_random(keys);

  // Create main socket
  relay_saddr.sin_family = AF_INET;
  relay_saddr.sin_port = PORT_6a44;
  sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

  // Create TUN device
  struct ifreq ifr;
  tun = open("/dev/net/tun", O_RDWR);
  set_zero(ifr);
  ifr.ifr_flags = IFF_TUN | IFF_NO_PI;
  strncpy(ifr.ifr_name, args::interface, IFNAMSIZ);
  ioctl(tun, TUNSETIFF, &ifr);

  // Set TUN device MTU
  ifr.ifr_mtu = 1280;
  ioctl(sock, SIOCSIFMTU, &ifr);

  // Bring TUN device up
  ioctl(sock, SIOCGIFFLAGS, &ifr);
  ifr.ifr_flags |= IFF_UP;
  ifr.ifr_flags &= ~IFF_MULTICAST;
  ioctl(sock, SIOCSIFFLAGS, &ifr);

  // Run interface up command
  if (args::interface_up) system(args::interface_up);

  // Bind main socket
  inet_pton(AF_INET, args::ipv4, &anycast_address);
  relay_saddr.sin_addr.s_addr = anycast_address;
  bind(sock, sa_cast(&relay_saddr), sizeof(relay_saddr));

  // Create and bind remote client receive socket
  if (args::remote_ipv4) {
    inet_pton(AF_INET, args::remote_ipv4, &relay_saddr.sin_addr);
    remote_sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    bind(remote_sock, sa_cast(&relay_saddr), sizeof(relay_saddr));
  }

  // Fork into background and write pid file
  if (!args::no_fork) {
    int pid = fork();
    if (pid > 0) {
      if (args::pidfile) std::ofstream(args::pidfile) << pid << endl;
      return 0;
    }
  }

  // Close standard file descriptors
  close(0);
  close(1);
  close(2);

  // Start work threads
  if (remote_sock) new std::thread(recv_ipv4_direct);
  new std::thread (recv_ipv6);
  new std::thread (recv_ipv4);

  // Periodacally re-key
  while (true) {
    sleep(600);
    bool ck = !current_key;
    set_random(keys[ck]);
    current_key = ck;
  }
}
